package com.example.yilanchatserver.application.controller;

import cn.hutool.core.util.ObjectUtil;
import com.alibaba.fastjson.JSON;
import com.example.yilanchatserver.business.freechat.*;
import com.example.yilanchatserver.business.freechat.service.RedisService;
import com.example.yilanchatserver.business.freechat.service.ErnieBotService;
import com.example.yilanchatserver.business.freechat.service.dto.FreeChatDTO;
import com.example.yilanchatserver.business.freechat.util.ErnieBotUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import javax.annotation.Resource;
import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.TimeUnit;

import static com.example.yilanchatserver.business.freechat.enums.CommonConstants.ERNIEBOT_ACCESS_TOKEN_KEY;


/**
 * 百度文心一言
 *
 * @author ziyu
 * @date 2023/9/5
 */
@ServerEndpoint(value = "/metah/freechat/erniebot/stream")
@Component
@Slf4j
public class ErnieBotWebSocketServer {
    private static int onlineCount = 0;
    // concurrent包的线程安全Set，用来存放每个客户端对应的MyWebSocket对象。
    private static CopyOnWriteArraySet<ErnieBotWebSocketServer> webSocketSet = new CopyOnWriteArraySet<>();

    private static ErnieBotService ernieBotService;

    //与某个客户端的连接会话，需要通过它来给客户端发送数据
    private Session session;

    private static RedisService redisService;

    @Resource
    private void setChatGPTService(ErnieBotService ernieBotService) {
        ErnieBotWebSocketServer.ernieBotService = ernieBotService;
    }

    @Resource
    private void setRedisService(RedisService redisService) {
        ErnieBotWebSocketServer.redisService = redisService;
    }

    private static String API_KEY = "6gbvl4VYDtm2OxDMPlTg5lnb";

    // 用户SECRET_KEY（请登录平台获取 https://cloud.baidu.com/ ）
    private static String SECRET_KEY = "UdTmakMcLjU63RsrLotMSU8s6Mcokw9N";

    private static String AUTH_TOKEN_URL = " https://aip.baidubce.com/oauth/2.0/token";

    private String getToken() throws Exception {
        try {
            String token = redisService.get(ERNIEBOT_ACCESS_TOKEN_KEY);
            if (StringUtils.isEmpty(token)) {
                token = ErnieBotUtils.getAccessToken(AUTH_TOKEN_URL, API_KEY, SECRET_KEY);
                redisService.set(ERNIEBOT_ACCESS_TOKEN_KEY, token, 8, TimeUnit.HOURS);
            }
            return token;
        } catch (Exception e) {
            log.error("获取token失败", e);
            throw e;
        }
    }

    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public void onOpen(Session session) {
        this.session = session;
        webSocketSet.add(this); //加入set中
        addOnlineCount(); //在线数加1
    }

    /**
     * 连接关闭调用的方法
     */
    @OnClose
    public void onClose() {
        webSocketSet.remove(this); //从set中删除
        subOnlineCount(); //在线数减1

    }

    /**
     * 收到客户端消息后调用的方法
     *
     * @param inputQuestion 客户端发送过来的消息
     */
    @OnMessage
    public void onMessage(String inputQuestion, Session session) throws Exception {

        String authToken = getToken();
        if (ObjectUtil.isEmpty(authToken)) {
            log.error("获取鉴权token失败");
            throw new BusinessException(BusinessErrorCodes.DEFAULT_BUSINESS_ERROR);
        }

        for (ErnieBotWebSocketServer item : webSocketSet) {
            if (item.session != session) {
                continue;
            }
            FreeChatDTO freeChatDTO = JSON.parseObject(inputQuestion, FreeChatDTO.class);
            Flux<String> response = ernieBotService.send(freeChatDTO, authToken);

            response.subscribe(message -> {
                try {
                    if (!"[Done]".equals(message) && !"[Error]".equals(message)) {
                        item.sendMessage(BaseResponseEntity.SUCCESS(message));
                        TimeUnit.MILLISECONDS.sleep(20);
                    } else if ("[Done]".equals(message)) {
                        item.sendMessage(BaseResponseEntity.FINISH());
                        item.closeSession();
                    } else {
                        item.sendMessage(BaseResponseEntity.ERROR());
                        item.closeSession();
                    }
                } catch (Exception e) {
                    try {
                        item.closeSession();
                    } catch (IOException ex) {
                        throw new RuntimeException(ex);
                    }
                }
            });
        }
    }

    @OnError
    public void onError(Session session, Throwable error) {
        log.error(error.getMessage());
    }

    public void closeSession() throws IOException {
        this.session.close();
    }

    public void sendMessage(BaseResponseEntity baseResponseEntity) throws IOException {
        String message = JSON.toJSONString(baseResponseEntity);
        if (this.session.isOpen()) {
            this.session.getBasicRemote().sendText(message);
        } else {
            webSocketSet.remove(this);
        }
    }

    public static synchronized int getOnlineCount() {
        return onlineCount;
    }

    public static synchronized void addOnlineCount() {
        ErnieBotWebSocketServer.onlineCount++;
    }

    public static synchronized void subOnlineCount() {
        ErnieBotWebSocketServer.onlineCount--;
    }
}

